package mashibing.class14;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * 并查集
 *
 * @author yujiafan
 * @version 1.0.0
 * @create 2023-06-01
 */
public class Class14_5_UnionFind {

    public static class Node<V> {
        V value;

        public Node(V v) {
            this.value = v;
        }
    }

    public static class UnionFind<V> {
        public HashMap<V, Node<V>> nodes;
        public HashMap<Node<V>, Node<V>> parents;
        public HashMap<Node<V>, Integer> sizeMap;

        public UnionFind(List<V> values) {
            nodes = new HashMap<>();
            parents = new HashMap<>();
            sizeMap = new HashMap<>();
            for (V value : values) {
                Node<V> node = new Node<>(value);
                nodes.put(value, node);
                parents.put(node, node);
                sizeMap.put(node, 1);
            }
        }

        // 给你一个节点，请你往上到不能再往上，把代表返回
        public Node<V> findFather(Node<V> cur) {
            Stack<Node<V>> stack = new Stack<>();
            while (cur != parents.get(cur)) {
                stack.add(cur);
                cur = parents.get(cur);
            }
            while (!stack.isEmpty()) {
                parents.put(stack.pop(), cur);
            }
            return cur;
        }

        public boolean isSameSet(V a, V b) {
            return findFather(nodes.get(a)) == findFather(nodes.get(b));
        }

        public void union(V a, V b) {
            Node<V> fatherA = findFather(nodes.get(a));
            Node<V> fatherB = findFather(nodes.get(b));
            if(fatherA != fatherB) {
                int sizeA = sizeMap.get(fatherA);
                int sizeB = sizeMap.get(fatherB);
                Node<V> big = sizeA >= sizeB ? fatherA : fatherB;
                Node<V> small = big == fatherA ? fatherB : fatherA;
                parents.put(small, big);
                sizeMap.remove(small);
                sizeMap.put(big, sizeA + sizeB);
            }
        }
    }
}
